{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset_path = \"Qwen2-Magpie-Pro-1M-v0.1.jsonl\" # Put your path to the dataset here\n",
    "dataset = load_dataset(\"json\", data_files=dataset_path)\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def high_quality_filter(example):\n",
    "    return (\n",
    "        example['input_quality'] in ['good', 'excellent']\n",
    "        and example['instruct_reward'] > -10\n",
    "        and not example['instruction'].endswith(':')\n",
    "        and (\n",
    "            example['min_similar_conversation_id'] is None\n",
    "            or example['conversation_id'] == example['min_similar_conversation_id']\n",
    "        )\n",
    "    )\n",
    "filtered_dataset = dataset['train'].filter(high_quality_filter)\n",
    "print(filtered_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_output_length(example):\n",
    "    return len(example['response'])\n",
    "\n",
    "filtered_dataset = filtered_dataset.map(lambda x: {'response_length': get_output_length(x)})\n",
    "# sort the dataset by response_length\n",
    "sorted_dataset = filtered_dataset.sort('response_length', reverse=True)\n",
    "# select the top 300000 examples\n",
    "top_x_dataset = sorted_dataset.select(range(300000))\n",
    "# Shuffle the dataset\n",
    "top_x_dataset = top_x_dataset.shuffle(seed=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit the dataset to the hub\n",
    "top_x_dataset.push_to_hub(\"Magpie-Align/Magpie-Qwen-Pro-300K-Filtered\", private=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "magpie",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
